#pragma once
#include <iostream>
#include <assert.h>
using namespace std;

enum Color{
  RED,
  BLACK
};

template <class T>
struct RBTreeNode{
  RBTreeNode<T> *_left;
  RBTreeNode<T> *_right;
  RBTreeNode<T> *_parent;

  T _data;
  Color _color;
  RBTreeNode(const T &data = T(), Color color = RED)
    :_left(nullptr),
    _right(nullptr),
    _parent(nullptr),
    _data(data),
    _color(color)
  {}
};

template<class T, class Ref, class Ptr>
class RBT_iterator{
  typedef RBT_iterator<T, Ref, Ptr> iterator;
  typedef RBTreeNode<T> Node;
  Node *_pnode;
public:
  RBT_iterator(Node *pnode = nullptr)
  :_pnode(pnode)
  {}

  Ref operator*() const{
    return _pnode->_data;
  }

  Ptr operator->() const{
    return &_pnode->_data;
  }
  
  bool operator==(const iterator &it) const{
    return _pnode == it._pnode;
  }

  bool operator!=(const iterator &it) const{
    return _pnode != it._pnode;
  }

  iterator& operator++(){
    if(_pnode->_right != nullptr)
    {
      Node *left = _pnode->_right;
      while(left->_left != nullptr)
      {
        left= left->_left;
      }
      _pnode = left;
    }
    else{
      Node *parent = _pnode->_parent;
      Node *cur = _pnode;
      while(parent != nullptr && cur == parent->_right)
      {
        parent = parent->_parent;
        cur = cur->_parent;
      }
      _pnode = parent;
    }
    return *this;
  }

  iterator& operator--(){
    if(_pnode->_left != nullptr)
    {
      Node *right = _pnode->_left;
      while(right->_right != nullptr)
      {
        right = right->_right;
      }
      _pnode = right;
    }
    else{
      Node *parent = _pnode->_parent;
      Node *cur = _pnode;
      while(parent != nullptr && cur == parent->_left)
      {
        parent = parent->_parent;
        cur = cur->_parent;
      }
      _pnode = parent;
    }
    return *this;
  }
};

template <class K, class T, class KofT>
class RBTree{
  typedef RBTreeNode<T> Node;
  Node *_root = nullptr;
public:
  typedef RBT_iterator<T, T&, T*> iterator;
  typedef RBT_iterator<T, const T&, const T*> const_iterator;

  iterator begin(){
    Node *left = _root;
    while(left != nullptr && left->_left != nullptr)
    {
      left = left->_left;
    }
    return iterator(left);
  }

  iterator end(){
    return iterator();
  }


  pair<iterator, bool> Insert(const T &data);

  iterator Find(const K &k);

  void Inorder(){
    _Inorder(_root);
    cout << endl;
  }

  int Height(){
    return _Height(_root);
  }

  bool Isbalance(){
  if(_root == nullptr) return true;
  if(_root->_color != BLACK)
  {
    cout << "违反性质2：根节点不为黑色！" << endl;
    return false;
  }

  int benchmark = 0;
  return PrevCheck(_root, 0, benchmark);
  }
private:
  void RotateL(Node *parent);
  void RotateR(Node *parent);
  void _Inorder(Node *root);
  int _Height(Node *root);
  bool PrevCheck(Node *root, int blacknum, int &benchmark);
};

template <class K, class T, class KofT>
bool RBTree<K,T,KofT>::PrevCheck(Node *root, int blacknum, int &benchmark){
  if(root == nullptr)
  {
    if(benchmark == 0)
    {
      benchmark = blacknum;
      return true;
    }
    else{
      if(blacknum != benchmark)
      {
        cout << "违反性质4：从任意节点到每个叶子节点的所有路径都包含相同数目的黑色节点！" << endl;
        return false;
      }
      else{
        return true;
      }
    }
  }

  if(root->_color == RED && root->_parent->_color == RED)
  {
    cout << "违反性质3：路径上有两个连续的红色节点！" << endl;
    return false;
  }

  if(root->_color == BLACK)
  {
    ++blacknum; 
  }
  return PrevCheck(root->_left, blacknum, benchmark)
      && PrevCheck(root->_right, blacknum, benchmark);
}

template <class K, class T, class KofT>
int RBTree<K,T,KofT>::_Height(Node *root){
  if(root == nullptr)
  {
    return 0;
  }
  int lefth = _Height(root->_left);
  int righth = _Height(root->_right);
  return lefth > righth? lefth+1 : righth+1;
}

template <class K, class T, class KofT>
void RBTree<K,T,KofT>::_Inorder(Node *root){
  KofT kot;
  if(root == nullptr) return; 
  _Inorder(root->_left);
  cout << kot(root->_data) << " ";
  _Inorder(root->_right);
}

template <class K, class T, class KofT>
pair<typename RBTree<K,T,KofT>::iterator, bool> RBTree<K,T,KofT>::Insert(const T &data)
{
  KofT kot;
  if(_root == nullptr)
  {
    _root = new Node(data, BLACK);
    return make_pair(iterator(_root), true);
  }

  Node *cur = _root;
  Node *parent = nullptr;
  while(cur != nullptr)
  {
    if(kot(data) > kot(cur->_data))
    {
      parent = cur;
      cur = cur->_right;
    }
    else if(kot(data) < kot(cur->_data))
    {
      parent  = cur;
      cur = cur->_left;
    }
    else{
      return make_pair(iterator(cur), true);
    }
  }
  cur = new Node(data,RED);
  Node *newnode = cur;
  if(kot(data) > kot(parent->_data))
  {
    parent->_right = cur;
  }
  else{
    parent->_left = cur;
  }
  cur->_parent = parent;
  
  while(parent != nullptr && parent->_color == RED) //上一次循环中grandparent 为根节点，此次循环parent == nullptr
  {
    Node *grandparent = parent->_parent;
    assert(grandparent != nullptr);
    assert(grandparent->_color == BLACK);

    Node *uncle = grandparent->_left;
    if(parent == grandparent->_left)
      uncle = grandparent->_right;

    if(uncle != nullptr && uncle->_color == RED) //情况一：uncle存在且为红
    {
      parent->_color = uncle->_color = BLACK;
      grandparent->_color = RED;
      cur = grandparent;
      parent = cur->_parent;
    }
    else //情况二、三：uncle不存在或uncle存在且为黑
    {
      if(parent == grandparent->_left)
      {
        if(cur == parent->_left) //左左
        {
          RotateR(grandparent);
          parent->_color = BLACK;
          grandparent->_color = RED;
        }
        else{ //左右
          RotateL(parent);
          RotateR(grandparent);
          cur->_color = BLACK;
          grandparent->_color = RED;
        }
      }
      else{
        if(cur == parent->_right) //右右
        {
          RotateL(grandparent);
          parent->_color = BLACK;
          grandparent->_color = RED;
        }
        else{ //右左
          RotateR(parent);
          RotateL(grandparent);
          cur->_color = BLACK;
          grandparent->_color = RED;
        }
      }
      break;
    } //end of else
  } //end of while
  
  if(cur == _root)
    cur->_color = BLACK;

  return make_pair(iterator(newnode), true);
}

template <class K, class T, class KofT>
void RBTree<K,T,KofT>::RotateL(Node *parent){
  Node *subR = parent->_right;
  Node *subRL = subR->_left;
  Node *ppNode = parent->_parent;

  parent->_right = subRL;
  if(subRL != nullptr)
  {
    subRL->_parent = parent;
  }

  subR->_left = parent;
  parent->_parent = subR;
 
  if(_root == parent)
  {
    _root = subR;
  }
  else{
    if(ppNode->_left == parent)
    {
      ppNode->_left = subR;
    }
    else{
      ppNode->_right = subR;
    }
  }
  subR->_parent = ppNode;
}

template <class K, class T, class KofT>
void RBTree<K,T,KofT>::RotateR(Node *parent){
  Node *subL = parent->_left;
  Node *subLR = subL->_right;
  Node *ppNode = parent->_parent;

  parent->_parent = subL;
  subL->_right = parent;

  parent->_left = subLR;
  if(subLR != nullptr)
  subLR->_parent = parent;

  if(_root == parent)
  {
    _root = subL;
  }
  else{
    if(ppNode->_left == parent)
    {
      ppNode->_left = subL;
    }
    else{
      ppNode->_right = subL;
    }
  }
  subL->_parent = ppNode;
}

template <class K, class T, class KofT>
typename RBTree<K,T,KofT>::iterator RBTree<K,T,KofT>::Find(const K &k){
  KofT kot;
  if(_root == nullptr)
  {
    return end();
  }
  Node *cur = _root;  
  while(cur != nullptr)
  {
    if(k > kot(cur->_data))
    {
      cur = cur->_right;
    }
    else if(k < kot(cur->_data))
    {
      cur = cur->_left;
    }
    else{
      return cur;
    }
  }
  return end();
}
